package com.dc.ai.repository;

import com.dc.ai.security.utils.SecurityUtils;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.concurrent.TimeUnit;

@Slf4j
@Component
@RequiredArgsConstructor
public class RedisChatHistoryRepository implements ChatHistoryRepository {

    private static final String HISTORY_KEY = "chat:history";
    private static final String USER_HISTORY_KEY = "chat:history:user";
    private static final String USER_CHAT_KEY = "chat:history:user:";
    private static final String MEMORY_KEY_PREFIX = "chat:memory:";
    private static final long EXPIRATION_DAYS = 30;

    private final RedisTemplate<String, Message> redisTemplate;
    private final StringRedisTemplate stringRedisTemplate;
    private final ObjectMapper redisObjectMapper;
    private final ChatMemory chatMemory;

    @Override
    public void save(String type, String chatId) {
        try {
            // 获取当前用户ID
            String userId = SecurityUtils.getUserId();

            updateChatIdsForType(type, chatId);

            // 如果有用户ID，更新用户相关的映射
            if (userId != null && !userId.isEmpty()) {
                // 保存聊天ID 用户ID映射
                stringRedisTemplate.opsForHash().put(USER_HISTORY_KEY, chatId, userId);
                stringRedisTemplate.expire(USER_HISTORY_KEY, EXPIRATION_DAYS, TimeUnit.DAYS);

                // 保存用户ID 聊天ID列表映射
                updateUserChatIds(userId, chatId);
            }
        } catch (Exception e) {
            log.error("保存聊天历史失败", e);
        }
    }

    private void updateChatIdsForType(String type, String chatId) throws JsonProcessingException {
        // 使用Redis事务减少并发问题
        Set<String> chatIds = new HashSet<>(getChatIds(type));
        chatIds.add(chatId);

        stringRedisTemplate.opsForHash().put(
                HISTORY_KEY,
                type,
                redisObjectMapper.writeValueAsString(chatIds)
        );
        stringRedisTemplate.expire(HISTORY_KEY, EXPIRATION_DAYS, TimeUnit.DAYS);
    }

    private void updateUserChatIds(String userId, String chatId) throws JsonProcessingException {
        String userChatKey = USER_CHAT_KEY + userId;
        String userChatIdsJson = stringRedisTemplate.opsForValue().get(userChatKey);

        Set<String> userChatIds = new HashSet<>();
        if (userChatIdsJson != null) {
            userChatIds = redisObjectMapper.readValue(userChatIdsJson, new TypeReference<HashSet<String>>() {
            });
        }

        userChatIds.add(chatId);
        stringRedisTemplate.opsForValue().set(
                userChatKey,
                redisObjectMapper.writeValueAsString(userChatIds),
                EXPIRATION_DAYS,
                TimeUnit.DAYS
        );
    }

    @Override
    public List<String> getChatIds(String type) {
        try {
            String json = (String) stringRedisTemplate.opsForHash().get(HISTORY_KEY, type);
            if (json == null) {
                return Collections.emptyList();
            }

            Set<String> readValue = redisObjectMapper.readValue(json, new TypeReference<HashSet<String>>() {
            });
            return new ArrayList<>(readValue);
        } catch (JsonProcessingException e) {
            log.error("反序列化聊天历史ID失败", e);
            return Collections.emptyList();
        }
    }

    /**
     * 根据用户ID获取聊天历史ID列表
     *
     * @param type   业务类型，如：chat、service、pdf
     * @param userId 用户ID
     * @return 聊天历史ID列表
     */
    @Override
    public List<String> getChatIdsByUserId(String type, String userId) {
        if (userId == null || userId.isEmpty()) {
            return Collections.emptyList();
        }

        try {
            // 直接从用户的聊天列表中获取
            String userChatKey = USER_CHAT_KEY + userId;
            String userChatIdsJson = stringRedisTemplate.opsForValue().get(userChatKey);

            if (userChatIdsJson == null) {
                return Collections.emptyList();
            }

            Set<String> allUserChatIds = redisObjectMapper.readValue(
                    userChatIdsJson,
                    new TypeReference<HashSet<String>>() {
                    }
            );

            // 获取指定类型的所有聊天ID
            List<String> typeChatIds = getChatIds(type);

            // 取交集：用户的聊天ID与指定类型的聊天ID
            allUserChatIds.retainAll(new HashSet<>(typeChatIds));

            return new ArrayList<>(allUserChatIds);
        } catch (JsonProcessingException e) {
            log.error("获取用户聊天历史失败", e);
            return Collections.emptyList();
        }
    }

    @PostConstruct
    private void init() {
        try {
            // 从Redis加载对话记忆
            loadChatMemoryData();
            log.info("成功从Redis加载聊天历史和记忆数据");
        } catch (Exception e) {
            log.error("从Redis加载聊天历史和记忆数据失败", e);
        }
    }

    private void loadChatMemoryData() {
        Set<String> keys = stringRedisTemplate.keys(MEMORY_KEY_PREFIX + "*");
        if (keys == null || keys.isEmpty()) {
            log.info("Redis中没有聊天记忆数据");
            return;
        }

        for (String key : keys) {
            try {
                String chatId = key.substring(MEMORY_KEY_PREFIX.length());
                List<Message> messages = redisTemplate.opsForList().range(key, 0, -1);

                if (messages != null && !messages.isEmpty()) {
                    // 获取聊天关联的用户ID
                    String userId = (String) stringRedisTemplate.opsForHash().get(USER_HISTORY_KEY, chatId);

                    // 如果有用户ID，将其添加到消息元数据中
                    if (userId != null && !userId.isEmpty()) {
                        for (Message message : messages) {
                            message.getMetadata().put("userId", userId);
                        }
                    }

                    chatMemory.add(chatId, messages);
                }
            } catch (Exception e) {
                log.error("加载聊天记忆数据失败: " + key, e);
                // 继续处理其他key，不中断整个加载过程
            }
        }
    }

    @PreDestroy
    private void persistent() {
        // Redis数据自动持久化，无需手动操作
        log.info("Chat history metadata and memory data persisted to Redis");
    }
}